import numpy as np
import pandas as pd
import logging
import time
from typing import Dict, List, Tuple, Optional, Callable

# 条件导入numba
try:
    from numba import jit, prange
    HAS_NUMBA = True
except ImportError:
    HAS_NUMBA = False
    # 定义空装饰器
    def jit(*args, **kwargs):
        def decorator(func):
            return func
        return decorator
    prange = range
    logging.warning("Numba未安装，字节级比对将使用纯NumPy实现（性能会有所下降）")

# 条件导入numexpr
try:
    import numexpr as ne
    HAS_NUMEXPR = True
except ImportError:
    HAS_NUMEXPR = False
    ne = None

from concurrent.futures import ThreadPoolExecutor, as_completed
from functools import lru_cache
from logger import logger


class DetailedFrameComparer:
    """详细帧比对器，支持精确的字节级比对
    
    使用NumPy向量化和Numba JIT编译加速比对操作
    """
    
    def __init__(self):
        self.logger = logging.getLogger('frame_analyzer')
        
    @staticmethod
    def _compare_bytes_numpy(arr1: np.ndarray, arr2: np.ndarray) -> Tuple[np.ndarray, int]:
        """使用纯NumPy的字节比对"""
        diff_mask = arr1 != arr2
        diff_indices = np.where(diff_mask)[0]
        return diff_indices, len(diff_indices)
        
    @staticmethod
    @jit(nopython=True, parallel=True)
    def _compare_bytes_numba(arr1: np.ndarray, arr2: np.ndarray) -> Tuple[np.ndarray, int]:
        """使用Numba加速的字节比对
        
        Returns:
            diff_indices: 不同字节的索引数组
            diff_count: 差异数量
        """
        n = len(arr1)
        diff_mask = np.zeros(n, dtype=np.bool_)
        
        for i in prange(n):
            if arr1[i] != arr2[i]:
                diff_mask[i] = True
                
        diff_indices = np.where(diff_mask)[0]
        return diff_indices, len(diff_indices)
        
    def _compare_bytes(self, arr1: np.ndarray, arr2: np.ndarray) -> Tuple[np.ndarray, int]:
        """智能选择比对方法"""
        if HAS_NUMBA:
            return self._compare_bytes_numba(arr1, arr2)
        else:
            return self._compare_bytes_numpy(arr1, arr2)
        
    def compare_payloads_detailed(self, 
                                 frame1: dict, 
                                 frame2: dict,
                                 compare_range: Optional[Tuple[int, int]] = None,
                                 max_diffs: int = 100) -> List[Dict]:
        """详细比对两个帧的payload
        
        Args:
            frame1: 基准帧
            frame2: 目标帧
            compare_range: 比对范围
            max_diffs: 最大报告差异数
            
        Returns:
            差异列表
        """
        try:
            # 提取payload数据
            payload1 = frame1['payload']
            payload2 = frame2['payload']
            
            # 获取帧格式信息
            format1 = frame1.get('frame_format', '1092')
            format2 = frame2.get('frame_format', '1092')
            
            # 根据格式限制payload长度
            if format1 == '896':
                payload1 = payload1[:896]
            elif format1 == '1024':
                payload1 = payload1[:1024]
                
            if format2 == '896':
                payload2 = payload2[:896]
            elif format2 == '1024':
                payload2 = payload2[:1024]
            
            # 应用比对范围
            if compare_range:
                start, end = compare_range
                # 确保不超出payload长度
                start = max(0, start)
                end1 = min(end, len(payload1) - 1)
                end2 = min(end, len(payload2) - 1)
                payload1 = payload1[start:end1+1]
                payload2 = payload2[start:end2+1]
                
            # 转换为NumPy数组
            arr1 = np.frombuffer(payload1, dtype=np.uint8)
            arr2 = np.frombuffer(payload2, dtype=np.uint8)
            
            # 长度检查
            if len(arr1) != len(arr2):
                return [{
                    'type': 'length_mismatch',
                    'length1': len(arr1),
                    'length2': len(arr2)
                }]
                
            # 使用Numba加速比对
            diff_indices, diff_count = self._compare_bytes(arr1, arr2)
            
            # 如果差异过多
            if diff_count > max_diffs:
                return [{
                    'type': 'too_many_diffs',
                    'total_diffs': diff_count,
                    'limit': max_diffs
                }]
                
            # 构建详细差异列表
            differences = []
            for idx in diff_indices[:max_diffs]:
                actual_position = int(idx)
                if compare_range:
                    actual_position += compare_range[0]
                    
                differences.append({
                    'byte_position': actual_position,
                    'value1': int(arr1[idx]),
                    'value2': int(arr2[idx]),
                    'diff_bits': bin(int(arr1[idx]) ^ int(arr2[idx]))[2:].zfill(8)
                })
                
            return differences
            
        except Exception as e:
            self.logger.error(f"比对payload失败: {str(e)}")
            return [{
                'type': 'error',
                'message': str(e)
            }]
        
    def batch_compare_frames(self,
                           frames_dict1: Dict[int, dict],
                           frames_dict2: Dict[int, dict],
                           frame_counts: List[int],
                           compare_range: Optional[Tuple[int, int]] = None) -> Dict:
        """批量比对帧数据
        
        使用向量化操作批量处理多个帧的比对
        """
        results = {
            'matching': [],
            'mismatching': [],
            'details': {}
        }
        
        # 准备批量数据
        batch_size = min(1000, len(frame_counts))
        
        for i in range(0, len(frame_counts), batch_size):
            batch_counts = frame_counts[i:i+batch_size]
            
            # 提取批量payload数据
            payloads1 = []
            payloads2 = []
            valid_counts = []
            
            for count in batch_counts:
                if count in frames_dict1 and count in frames_dict2:
                    frame1 = frames_dict1[count]
                    frame2 = frames_dict2[count]
                    
                    # 提取比对数据
                    payload1 = frame1['payload']
                    payload2 = frame2['payload']
                    
                    if compare_range:
                        start, end = compare_range
                        payload1 = payload1[start:end+1]
                        payload2 = payload2[start:end+1]
                        
                    payloads1.append(payload1)
                    payloads2.append(payload2)
                    valid_counts.append(count)
            
            # 批量比较
            if valid_counts:
                # 使用NumPy批量比较
                match_results = self._batch_compare_payloads(payloads1, payloads2)
                
                # 记录结果
                for idx, (count, is_match) in enumerate(zip(valid_counts, match_results)):
                    if is_match:
                        results['matching'].append(count)
                    else:
                        results['mismatching'].append(count)
                        # 详细比对不匹配的帧
                        frame1 = frames_dict1[count]
                        frame2 = frames_dict2[count]
                        details = self.compare_payloads_detailed(
                            frame1, frame2, compare_range
                        )
                        if details:
                            results['details'][count] = details
                            
        return results
        
    def _batch_compare_payloads(self, payloads1: List[bytes], payloads2: List[bytes]) -> List[bool]:
        """批量比较payload数据"""
        match_results = []
        
        for p1, p2 in zip(payloads1, payloads2):
            # 快速比较
            if len(p1) != len(p2):
                match_results.append(False)
            else:
                # 使用NumPy数组比较
                arr1 = np.frombuffer(p1, dtype=np.uint8)
                arr2 = np.frombuffer(p2, dtype=np.uint8)
                match_results.append(np.array_equal(arr1, arr2))
                
        return match_results


class OptimizedVirtualChannelComparer:
    """优化的虚拟信道比对器
    
    使用高效的数据结构和算法进行虚拟信道级别的比对
    """
    
    def __init__(self):
        self.logger = logging.getLogger('frame_analyzer')
        self.detailed_comparer = DetailedFrameComparer()
        
    def compare_virtual_channel_optimized(self,
                                        vc_frames1: List[dict],
                                        vc_frames2: List[dict],
                                        vc_id: int,
                                        compare_range: Optional[Tuple[int, int]] = None) -> Dict:
        """优化的虚拟信道比对
        
        使用高效的算法和数据结构
        """
        # 创建帧计数到帧的映射（使用字典以提高查找效率）
        frames1_dict = {int(f['frame_count']): f for f in vc_frames1}
        frames2_dict = {int(f['frame_count']): f for f in vc_frames2}
        
        # 使用NumPy数组进行集合运算
        counts1 = np.array(list(frames1_dict.keys()), dtype=np.int64)
        counts2 = np.array(list(frames2_dict.keys()), dtype=np.int64)
        
        # 高效的集合运算
        common_counts = np.intersect1d(counts1, counts2, assume_unique=True)
        missing_in_target = np.setdiff1d(counts1, counts2, assume_unique=True)
        missing_in_base = np.setdiff1d(counts2, counts1, assume_unique=True)
        
        # 初始化结果
        result = {
            'vc_id': vc_id,
            'frames_in_base': len(counts1),
            'frames_in_target': len(counts2),
            'matching_frames': 0,
            'mismatching_frames': 0,
            'missing_in_target': len(missing_in_target),
            'missing_in_base': len(missing_in_base),
            'frame_comparisons': []
        }
        
        # 批量比对共同帧
        if len(common_counts) > 0:
            batch_result = self.detailed_comparer.batch_compare_frames(
                frames1_dict, frames2_dict, common_counts.tolist(), compare_range
            )
            
            result['matching_frames'] = len(batch_result['matching'])
            result['mismatching_frames'] = len(batch_result['mismatching'])
            
            # 构建详细结果
            # 匹配的帧
            for count in batch_result['matching']:
                result['frame_comparisons'].append({
                    'frame_count': int(count),
                    'type': 'matching',
                    'vc_id': vc_id
                })
                
            # 不匹配的帧
            for count in batch_result['mismatching']:
                frame_comparison = {
                    'frame_count': int(count),
                    'type': 'content_mismatch',
                    'vc_id': vc_id
                }
                
                # 添加详细差异
                if count in batch_result['details']:
                    frame_comparison['differences'] = [{
                        'field': 'payload',
                        'differences': batch_result['details'][count]
                    }]
                    
                result['frame_comparisons'].append(frame_comparison)
        
        # 添加缺失帧的记录
        for count in missing_in_target:
            result['frame_comparisons'].append({
                'frame_count': int(count),
                'type': 'missing_in_target',
                'vc_id': vc_id
            })
            
        for count in missing_in_base:
            result['frame_comparisons'].append({
                'frame_count': int(count),
                'type': 'missing_in_base',
                'vc_id': vc_id
            })
            
        return result


class MemoryEfficientComparer:
    """内存高效的比对器
    
    使用内存映射和流式处理减少内存占用
    """
    
    def __init__(self, cache_size: int = 1000):
        self.cache_size = cache_size
        self.logger = logging.getLogger('frame_analyzer')
        self._frame_cache = {}
        
    @lru_cache(maxsize=10000)
    def _compute_payload_hash(self, payload_bytes: bytes) -> int:
        """计算payload的哈希值（带缓存）"""
        return hash(payload_bytes)
        
    def compare_with_memory_mapping(self,
                                  frames1: List[dict],
                                  frames2: List[dict],
                                  compare_range: Optional[Tuple[int, int]] = None) -> Dict:
        """使用内存映射进行比对
        
        适用于大文件，减少内存占用
        """
        # 构建索引结构（只存储必要信息）
        index1 = self._build_frame_index(frames1)
        index2 = self._build_frame_index(frames2)
        
        # 使用生成器进行流式比对
        comparison_generator = self._streaming_compare(
            frames1, frames2, index1, index2, compare_range
        )
        
        # 收集结果
        result = {
            'total_frames': [len(frames1), len(frames2)],
            'matching_frames': 0,
            'virtual_channel_results': {}
        }
        
        # 流式处理结果
        for vc_id, vc_result in comparison_generator:
            result['virtual_channel_results'][vc_id] = vc_result
            result['matching_frames'] += vc_result['matching_frames']
            
            # 定期清理缓存
            if len(self._frame_cache) > self.cache_size:
                self._frame_cache.clear()
                
        return result
        
    def _build_frame_index(self, frames: List[dict]) -> pd.DataFrame:
        """构建轻量级帧索引"""
        data = []
        for idx, frame in enumerate(frames):
            payload = frame['payload']
            if isinstance(payload, (bytes, bytearray)):
                payload_hash = self._compute_payload_hash(bytes(payload))
            else:
                payload_hash = 0
                
            data.append({
                'idx': idx,
                'frame_count': int(frame['frame_count']),
                'virtual_channel': int(frame.get('virtual_channel', 0)),
                'payload_hash': payload_hash,
                'payload_size': len(frame['payload'])
            })
            
        df = pd.DataFrame(data)
        df.set_index(['virtual_channel', 'frame_count'], inplace=True)
        return df
        
    def _streaming_compare(self, frames1, frames2, index1, index2, compare_range):
        """流式比对生成器"""
        # 获取所有虚拟信道
        all_vc_ids = sorted(set(
            index1.index.get_level_values('virtual_channel').unique().tolist() +
            index2.index.get_level_values('virtual_channel').unique().tolist()
        ))
        
        # 逐个处理虚拟信道
        for vc_id in all_vc_ids:
            # 获取该虚拟信道的索引
            vc_index1 = index1.xs(vc_id, level='virtual_channel') if vc_id in index1.index.get_level_values('virtual_channel') else None
            vc_index2 = index2.xs(vc_id, level='virtual_channel') if vc_id in index2.index.get_level_values('virtual_channel') else None
            
            # 比对该虚拟信道
            vc_result = self._compare_single_vc_streaming(
                vc_id, vc_index1, vc_index2, frames1, frames2, compare_range
            )
            
            yield vc_id, vc_result
            
    def _compare_single_vc_streaming(self, vc_id, vc_index1, vc_index2, 
                                   frames1, frames2, compare_range):
        """流式比对单个虚拟信道"""
        # 初始化结果
        result = {
            'vc_id': vc_id,
            'frames_in_base': len(vc_index1) if vc_index1 is not None else 0,
            'frames_in_target': len(vc_index2) if vc_index2 is not None else 0,
            'matching_frames': 0,
            'mismatching_frames': 0,
            'missing_in_target': 0,
            'missing_in_base': 0,
            'frame_comparisons': []
        }
        
        # 处理空虚拟信道
        if vc_index1 is None or vc_index2 is None:
            # 处理缺失情况
            if vc_index1 is None and vc_index2 is not None:
                result['missing_in_base'] = len(vc_index2)
            elif vc_index2 is None and vc_index1 is not None:
                result['missing_in_target'] = len(vc_index1)
            return result
            
        # 获取帧计数
        counts1 = set(vc_index1.index.tolist())
        counts2 = set(vc_index2.index.tolist())
        
        # 计算差集
        common_counts = counts1 & counts2
        missing_in_target = counts1 - counts2
        missing_in_base = counts2 - counts1
        
        result['missing_in_target'] = len(missing_in_target)
        result['missing_in_base'] = len(missing_in_base)
        
        # 比对共同帧（使用缓存的哈希值进行快速比对）
        for count in common_counts:
            row1 = vc_index1.loc[count]
            row2 = vc_index2.loc[count]
            
            # 首先比较哈希值
            if row1['payload_hash'] == row2['payload_hash'] and row1['payload_size'] == row2['payload_size']:
                result['matching_frames'] += 1
                result['frame_comparisons'].append({
                    'frame_count': int(count),
                    'type': 'matching',
                    'vc_id': vc_id
                })
            else:
                # 哈希值不同，需要详细比对
                result['mismatching_frames'] += 1
                result['frame_comparisons'].append({
                    'frame_count': int(count),
                    'type': 'content_mismatch',
                    'vc_id': vc_id,
                    'differences': [{'field': 'payload', 'differences': []}]
                })
                
        return result 